from typing import Annotated

from langchain_core.messages import ToolMessage
from langchain_core.runnables import RunnableConfig
from langchain_core.tools import tool, InjectedToolCallId
from langgraph.prebuilt import InjectedState
from langgraph.types import Command

from agent.my_state import CustomState


@tool
def get_user_name(
    tool_call_id: Annotated[str, InjectedToolCallId],
    config: RunnableConfig
) -> Command:
    """获取当前用户的username，以便生成祝福语句"""
    user_name = config['configurable'].get("user_name", 'gary')
    print(f"调用工具 get_user_name， 传入的用户名是: {user_name}")
    # 模拟
    return Command(update={
        "username": user_name,  # 更新状态中的用户名,
        # 更新一条工具执行后的消息： ToolMessage类型
        "messages": [
            ToolMessage(
                content="成功的得到当前用户的username",
                tool_call_id=tool_call_id
            )
        ]
    })


@tool
def greet_user(state: Annotated[CustomState, InjectedState]) -> str:
    """在获取用户的username之后，生成祝福语句"""
    username = state['username']  # 从状态中获取用户名
    print(f"从state中获取到的用户名并为其生成祝福语的用户是: {username}")
    return f'祝贺你：{username}！'


print(get_user_name.name)
print(get_user_name.description)
print(get_user_name.args)
print(get_user_name.args_schema.model_json_schema())
print(get_user_name.return_direct)

# print(get_user_name.invoke({'a': 40, 'b': 2, 'operation': 'multiply'}))

print(greet_user.name)
print(greet_user.description)
print(greet_user.args)
print(greet_user.args_schema.model_json_schema())
print(greet_user.return_direct)
